#pragma once
#ifndef TOFTLIB_MEMORY_POOL_H
#define TOFTLIB_MEMORY_POOL_H

namespace toftlib
{
    struct MemNode
    {
        size_t size;

        union
        {
            MemNode* freeListLink;
            char* address;
        };
    };

    template <class Locker>
    class MemPool
    {
    public:
        enum{MAX_INDEX = 16};
        enum{ALIGN = 2};


        MemPool()
        {
            allocSize = RoundUp(GetSystemPageSize(), ALIGN);
            chunkSize = 128;
            Initialize();
        }
        MemPool(size_t wantedMinAllocSize, size_t wantedChunkSize)
        {
            size_t systemPageSize = GetSystemPageSize();

            if(wantedMinAllocSize <= systemPageSize)
            {
                allocSize = RoundUp(systemPageSize, ALIGN);
            }
            else
            {
                allocSize = RoundUp(RoundUp(wantedMinAllocSize, systemPageSize), ALIGN);
            }

            chunkSize = RoundUp(wantedChunkSize, ALIGN);

            Initialize();
        }

        MemNode* Allocate(size_t n)
        {
            MemNode* result = NULL;

            n = RoundUp(n + sizeof(MemNode), chunkSize);

            locker.Lock();

            // If wanted size is bigger than maxFixedSize, use freeList[0]
            if(n > maxFixedSize)
            {
                n = RoundUp(n, allocSize);

                result = freeList[0];
                while(result != NULL  && n > result->size)
                {
                    result = (MemNode*)result->freeListLink;
                }

                // Find a position
                if(result != NULL)
                {
                    if(NULL == result->freeListLink)
                    {
                        freeList[0] = NULL;
                    }
                    else
                    {
                        freeList[0]->freeListLink = result->freeListLink;
                    }
                    result->address = (char*)result + sizeof(MemNode);

                    locker.Unlock();
                    return result;
                }

                locker.Unlock();
            }
            else
            {
                size_t pos = GetFreeListIndex(n, chunkSize);

                result = freeList[pos];
                while(result == NULL && pos < MAX_INDEX - 2)
                {
                    ++pos;    
                    result = freeList[pos];                                  
                }

                // Find a node or move to last postion of free list
                if(result != NULL)
                {
                    if(NULL == result->freeListLink)
                    {
                        freeList[pos] = NULL;
                    }
                    else
                    {
                        freeList[pos] = result->freeListLink;
                    }
                    result->address = (char*)result + sizeof(MemNode);

                    locker.Unlock();
                    return result;
                }

            }


            if(n >= allocSize)
            {
                locker.Unlock();
                result = (MemNode*) new char[n];
                result->address = (char*)result + sizeof(MemNode);
                result->size = n;
            }
            // Do not allocate too small memory
            else
            {
                char* memory = new char[allocSize];
                size_t copys = allocSize / n; 
                
                result = (MemNode*)memory;
                result->address = (char*)result + sizeof(MemNode);
                result->size = n;

                size_t pos = GetFreeListIndex(n, chunkSize);
                for(size_t i = 1; i <= copys - 1; ++i)
                {
                    MemNode* oneNode = (MemNode*)(memory + i * n);
                    oneNode->size = n;
                    oneNode->freeListLink = freeList[pos];
                    freeList[pos] = oneNode;
                    currentFreeMemory += oneNode->size;
                }

                locker.Unlock();
            }
            
            return result;
        }

        void Deallocate(MemNode* memory)
        {
            locker.Lock();
            if(memory->size > maxFixedSize)
            {
                memory->freeListLink = freeList[0];
                freeList[0] = memory;
                currentFreeMemory += memory->size;
            }
            else
            {
                size_t pos = GetFreeListIndex(memory->size, chunkSize);

                memory->freeListLink = freeList[pos];
                freeList[pos] = memory;
                currentFreeMemory += memory->size;
            }

            locker.Unlock();
        }

        void FreeMemory()
        {
            
            locker.Lock();
            // Free big memory at freeList[0] first
            MemNode* node = freeList[0];
            MemNode* tempNode = NULL;
            while(node != NULL)
            {
                tempNode = node;
                node = node->freeListLink;
                currentFreeMemory -= tempNode->size;
                delete tempNode;
            }

            //Free other slot which chunkSize bigger than allocSize
            size_t currentChunkSize = 0;
            for(size_t i = 1; i <= MAX_INDEX; ++i)
            {
                currentChunkSize += chunkSize;
                if(currentChunkSize < allocSize)
                {
                    continue;
                }
                
                node = freeList[i];
                while(node != NULL)
                {
                    tempNode = node;
                    node = node->freeListLink;
                    currentFreeMemory -= tempNode->size;
                    delete tempNode;

                }
            }

            locker.Unlock();
        }

    private:

        void Initialize()
        {
            for(int i = 0; i < MAX_INDEX; ++i)
            {
                 freeList[i] = NULL;
            }

            maxFixedSize = chunkSize * MAX_INDEX;
            currentFreeMemory = 0;
        }

        static size_t RoundUp(size_t bytes, size_t boundary)
        {
            return ((bytes + boundary - 1) & ~(boundary - 1));
        }

        static size_t GetFreeListIndex(size_t bytes, size_t boundary)
        {
            return (bytes + 1) / boundary;
        }

        static size_t GetSystemPageSize()
        {
            SYSTEM_INFO si;
            GetSystemInfo(&si);

            return si.dwPageSize;
        }

        MemNode* volatile freeList[MAX_INDEX];

        Locker locker;

        size_t allocSize;
        size_t chunkSize;
        size_t maxFixedSize;
        size_t currentFreeMemory;
    };

}


#endif
